package com.mini.ws;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import com.mini.config.MyWebSocketConfig;
import com.mini.constant.CacheConstants;
import com.mini.enums.YesOrNoEnum;
import com.mini.redis.RedisCache;
import com.mini.util.ChannelUtils;
import com.mini.util.SpringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Collection;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * 处理 WebSocket 连接
 *
 * @author click33
 * @since 2022-2-11
 */
@Slf4j
public class MyWebSocketHandler extends TextWebSocketHandler {


    private final MyWebSocketConfig config = SpringUtils.getBean(MyWebSocketConfig.class);

    private final RedisCache redisCache = SpringUtils.getBean(RedisCache.class);

    /**
     * 固定前缀
     */
    private static final String USER_ID = "user_id_";

    /**
     * 存放Session集合，方便推送消息
     */
    private static ConcurrentHashMap<String, Set<WebSocketSession>> webSocketSessionMaps = new ConcurrentHashMap<>();

    // 监听：连接开启 
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {

        // put到集合，方便后续操作
        String userId = session.getAttributes().get("userId").toString();

        Set<WebSocketSession> webSocketSessions = webSocketSessionMaps.get(USER_ID + userId);

        if (CollectionUtil.isEmpty(webSocketSessions)) {
            webSocketSessions = new HashSet<>();
        }

        webSocketSessions.add(session);

        webSocketSessionMaps.put(USER_ID + userId, webSocketSessions);

        // 给个提示
        String tips = "Web-Socket 连接成功，sid=" + session.getId() + "，userId=" + userId;

        System.out.println(tips);

        sendMessage(session, tips);

        setHeartBeat(session);
    }

    // 监听：连接关闭 
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        // 从集合移除
        String userId = session.getAttributes().get("userId").toString();

        Set<WebSocketSession> webSocketSessions = webSocketSessionMaps.get(USER_ID + userId);

        webSocketSessions.remove(session);

        log.info("Set size:{}", webSocketSessions.size());

        // 给个提示 
        String tips = "Web-Socket 连接关闭，sid=" + session.getId() + "，userId=" + userId;
        System.out.println(tips);
        ChannelUtils.cancelChannel(userId);
    }

    // 收到消息 
    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {

        if (StrUtil.equals("heartbeat", message.getPayload())) {
            System.out.println("sid为：" + session.getId() + "，心跳包：" + message);
            setHeartBeat(session);
        } else {
            System.out.println("sid为：" + session.getId() + "，发来：" + message);
        }

        sendMessage(session, message.toString());
    }

    // ----------- 

    // 向指定客户端推送消息 
    public static void sendMessage(WebSocketSession session, String message) {
        try {
            System.out.println("向sid为：" + session.getId() + "，发送：" + message);
            session.sendMessage(new TextMessage(message));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    // 向指定用户推送消息
    public static boolean sendMessage(long userId, String message) {

        Set<WebSocketSession> webSocketSessions = webSocketSessionMaps.get(USER_ID + userId);

        for (WebSocketSession session : webSocketSessions) {
            if (session != null) {
                sendMessage(session, message);
            }
        }

        return Boolean.TRUE;
    }

    public static Long getConnectionSize() {

        Collection<Set<WebSocketSession>> values = MyWebSocketHandler.webSocketSessionMaps.values();

        long size = 0;

        for (Set<WebSocketSession> value : values) {
            size += value.size();
        }

        return size;
    }

    public static void close(String userId, String sessionId) {

        Set<WebSocketSession> webSocketSessions = webSocketSessionMaps.get(USER_ID + userId);

        if (CollectionUtil.isEmpty(webSocketSessions)) {
            return;
        }

        log.info("心跳过期，关闭连接:sid={}, userId={}", sessionId, userId);

        for (WebSocketSession session : webSocketSessions) {
            if (StrUtil.equals(sessionId, session.getId())) {
                try {
                    sendMessage(session, "心跳过期,关闭连接!");
                    session.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    private void setHeartBeat(WebSocketSession session) {

        if (YesOrNoEnum.YES.equals(config.getOpenHeartBate())) {
            String userId = session.getAttributes().get("userId").toString();
            String key = StrUtil.format(CacheConstants.WEBSOCKET_HEART_BEAT, userId, session.getId());
            redisCache.setCacheObject(key, userId + ":" + session.getId(), config.getHeartBateTime(), TimeUnit.SECONDS);
        }
    }
}

